import brian_no_units
from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
from time import clock, time
defaultclock.dt = 0.1
#set_global_preferences(useweave=True,usecodegenweave = True)

# Conductances (mS/cm^2)
GNa = 120 
GNap = 0.1
GK_dr = 100 
GCa_NS = 14 
GCa_ND = .03 
GK_CaS = 3.136  # canonical: 5
GK_CaD = 0.69   # canonical: 1.1
GCa_L = 0.33 
gleak = 0.51 

# Static parameters
C = 1 
gc = 0.1        # coupling conductance (mS/cm^2)
p = 0.1
Kd = 0.2      # uM
f = 0.01      # percent free to bound Ca
alpha = 0.009 # mol/C/um
kca = 2       # Ca removal rate

# Half Activation voltages in mV, Slopes in MV, Time Constants in milliseconds
Vhm = -35
Sm = -7.8
Vhh = -55
Sh = 7
Vhmnap = -47.1
Smnap = -3.1
Vhhnap = -59
Shnap = 8
Vhn = -28
Sn = -15
VhmN = -30
SmN = -5
VhhN = -45
ShN = 5
VhmL = -40
SmL = -7
TaumN = 4 
TauhN = 40
TaumL = 40
Taum = 0.000001
Taumnap = Taum

# Reversal potentials in mV
ENa = 55 
EK = -80 
ECa = 80 
Eleak = -60 

def iext2(t):
    return 0

def iext(t):
    if t <=500:
        return -20
    elif t <= 5500:
        return -20+(t-500) * 25/5000
    elif t <= 100000:
        return (5-(t-5500) * 25/5000)
    else:
        return 0
def iext1(t):
    if t <= 500:
        return 0
    elif t <= 5500:
        return (t-500) * 25/5000
    elif t <= 100000:
        return (25-(t-5500) * 25/5000)
    else:
        return 0

eqsMN=Equations('''
dh/dt= (hinf-h)/Tauh : 1
dn/dt= (ninf-n)/Taun : 1
dmnS/dt = (mnSinf-mnS)/TaumN : 1
dhnS/dt = (hnSinf-hnS)/TauhN : 1
dmnD/dt = (mnDinf-mnD)/TaumN : 1
dhnD/dt = (hnDinf-hnD)/TauhN : 1
dm/dt = (minf-m)/Taum : 1
dmnap/dt = (mnapinf-mnap)/Taumnap : 1
dhnap/dt = (hnapinf-hnap)/Tauhnap : 1
dml/dt = (mlinf-ml)/TaumL : 1
dCaS/dt = f*(-alpha*ICaS-kca*CaS) : 1
dCaD/dt = f*(-alpha*ICaD-kca*CaD) : 1
dvm/dt = 1/C*(iext2(t)-INaS-IKS-ICaS-IleakS-IcouplingS) : 1
dVd/dt = 1/C*(-INapD-IKD-ICaD-IleakD-IcouplingD) : 1

Tauh = 30/(exp((vm+50)/15)+exp(-(vm+50)/16)) : 1
Taun = 7/(exp((vm+40)/40)+exp(-(vm+40)/50)) : 1
Tauhnap = 1200/(cosh(Vd + 59)/16) :1
minf = 1/(1+exp((vm-Vhm)/Sm)) : 1
hinf = 1/(1+exp((vm-Vhh)/Sh)) : 1
mnapinf = 1/(1+exp((Vd-Vhmnap)/Smnap)) : 1
hnapinf = 1/(1+exp((Vd-Vhhnap)/Shnap)) : 1
ninf = 1/(1+exp((vm-Vhn)/Sn)) : 1
mnSinf = 1/(1+exp((vm-VhmN)/SmN)) : 1
hnSinf = 1/(1+exp((vm-VhhN)/ShN)) : 1
mnDinf = 1/(1+exp((Vd-VhmN)/SmN)) : 1
hnDinf = 1/(1+exp((Vd-VhhN)/ShN)) : 1
mlinf = 1/(1+exp((Vd-VhmL)/SmL)) : 1
GNap : 1

INapD = GNap*mnap*hnap*(Vd-ENa) : 1
INaS = GNa*m**3*h*(vm-ENa) : 1
IKS = (GK_dr*n**4 + GK_CaS*CaS/(CaS+Kd))*(vm-EK) : 1
ICaS = GCa_NS*mnS**2*hnS* (vm-ECa) : 1
IleakS = gleak*(vm-Eleak) : 1
IcouplingS = gc/p*(vm-Vd) : 1
IKD = GK_CaD*CaD/(CaD+Kd)*(Vd-EK) : 1
ICaD = (GCa_ND*mnD**2*hnD+GCa_L*ml)*(Vd-ECa) : 1
IleakD = gleak*(Vd-Eleak) : 1
IcouplingD = gc/(1-p)*(Vd-vm) : 1
''')

P=NeuronGroup(240,model=eqsMN,
threshold=EmpiricalThreshold(threshold=-40),
implicit=True)
P2=NeuronGroup(80,model=eqsMN,
threshold=EmpiricalThreshold(threshold=-40),
implicit=True)
P1=NeuronGroup(160,model=eqsMN,
threshold=EmpiricalThreshold(threshold=-40),
implicit=True)


ivm= -60
ivd= -60
for p in [P,P1,P2]:
    p.vm= ivm
    p.Vd= ivd
    #P.h = 1/(1+exp((ivm-Vhh)/Sh))
    p.h = 0.9
    #P.n = 1/(1+exp((ivm-Vhn)/Sn))
    p.n = 0
    #P.mnS = 1/(1+exp((ivm-VhmN)/SmN))
    p.mnS = 0
    #P.hnS = 1/(1+exp((ivm-VhhN)/ShN))
    p.hnsS = 0.9
    #P.mnD = 1/(1+exp((ivd-VhmN)/SmN))
    p.mnD = 0
    #P.hnD = 1/(1+exp((ivd-VhhN)/ShN))
    p.hnD = 0.9
    #P.ml = 1/(1+exp((ivd-VhmL)/SmL))
    p.mnap = 0
    p.hnap = 0.9
    p.ml = 0
    p.m = 0
    p.CaS = 0
    p.CaD = 0
    p.GNap = 0.1
    


tracevm=StateMonitor(P,'vm',record=[0,1])
tracevm1=StateMonitor(P1,'vm',record=[0])
tracevm2=StateMonitor(P2,'vm',record=[0])
#traceVd=StateMonitor(P,'Vd',record=[0])
#traceINap=StateMonitor(P,'INapD',record=[0])
spikes = SpikeMonitor(P)

net1 = Network(P,tracevm)

net2 = Network(P1,P2,tracevm1,tracevm2)

DUR = 2000

start = clock()

net2.run(DUR)

tnet1 = clock()
print tnet1-start

net1.run(DUR)

tnet2 = clock()
print tnet2-tnet1



#
#figure(1)
#ion()
#net.run(1)
#subplot(211)
#tracevm.plot(0,refresh=100,showlast=DUR)
#subplot(212)
#tracevm.plot(1,refresh=100,showlast=DUR)
#net.run(DUR+1)
#ioff()

